require 'CSV'

# TemplateBuilder
# Create a published_vulnerability_attack template by parsing an engine_log.csv file
# The engine_log.csv file is found in the analysis export of a previous analysis run
# The template will contain all variants referenced by their CVE Id.
# The filter method may be updated to filter on other columns of the engine_log.csv
# An example of 'Service' is provided.
class PvaTemplateBuilder
  attr_accessor :csv
  attr_accessor :heading
  attr_accessor :template
  
  def initialize(filepath = "engine_log.csv")
    @filepath = filepath
  end
  
# parse the csv file. Linefeeds should be either UNIX or Windows style
# Any double quotes should be removed, and each column should have not ','s
  def parse
    @csv = CSV.parse(File.read(@filepath))
    @heading = Hash.new
    @csv[0].each_index do |h|
      @heading[@csv[0][h]] = h
    end
  end
  
  def filter
    # Assumimg this is an engine log, filter on the second entry for each variant
    @csv.delete_if { |row| !(row[@heading['Category']] == 'Variant' and row[@heading['Event']] == 'End') }
    # Additional Filters can go here
    # Filter by Service example:
    # @csv.delete_if { |row| !(row[@heading['Service']] == 'http') }
  end

  def create_template
    @template =  "<?xml version=\"1.0\" encoding=\"UTF-8\"?>
<mu_config version=\"3.0\">
<templates>
<published_network_attack>
  <analyzer_mode>ClientServer</analyzer_mode>
  <evasion>NONE</evasion>
  <filters>
    <reference_filter>
      <operator>IS</operator>
      <values>\n"
    @csv.each do |row|
      next if row[@heading['CVE']].nil?
      @template += "      <value>
          <reference>CVE</reference>
          <value>#{row[@heading['CVE']]}</value>
      </value>\n"
      end
    @template += "      </values>
    </reference_filter>
  </filters>
</published_network_attack>
</templates>
</mu_config>"
  end

end

# Filepath to an engine_log.csv file
filepath = (ARGV[0].nil?) ? 'engine_log.csv' : ARGV[0]

# Create our builder and parse the file
p = PvaTemplateBuilder.new(filepath)
p.parse

# Filter out engine log entries other than the Variant End events
# Comment this line out if you are using a simple list of CVE references - Where each row contains a valid CVE value
p.filter

# Buile xml template
p.create_template

# Print to STDOUT
puts p.template

# Write to a File
File.open("template.xml", 'w') {|f| f.write(p.template) }

